# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.
import math
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F

from fairseq import options, utils
from fairseq.models import (
    FairseqIncrementalDecoder,
    FairseqLanguageModel,
    register_model,
    register_model_architecture,
)
from fairseq.modules import (
    AdaptiveSoftmax,
    LayerNorm,
    MultiheadAttention,
    PositionalEmbedding,
    SinusoidalPositionalEmbedding,
    AdaptiveInput,
    CharacterTokenEmbedder,
)

@register_model('transformer_lm_bpe_coherence_disef')
class BPECDEFTransformerLanguageModel(FairseqLanguageModel):
    def __init__(self, decoder):
        super().__init__(decoder)

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        # fmt: off
        parser.add_argument('--activation-fn',
                            choices=utils.get_available_activation_fns(),
                            help='activation function to use')
        parser.add_argument('--dropout', default=0.1, type=float, metavar='D',
                            help='dropout probability')
        parser.add_argument('--attention-dropout', default=0., type=float, metavar='D',
                            help='dropout probability for attention weights')
        parser.add_argument('--relu-dropout', default=0., type=float, metavar='D',
                            help='dropout probability after ReLU in FFN')
        parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
                            help='decoder embedding dimension')
        parser.add_argument('--decoder-output-dim', type=int, metavar='N',
                            help='decoder output dimension')
        parser.add_argument('--decoder-input-dim', type=int, metavar='N',
                            help='decoder input dimension')
        parser.add_argument('--decoder-ffn-embed-dim', type=int, metavar='N',
                            help='decoder embedding dimension for FFN')
        parser.add_argument('--decoder-layers', type=int, metavar='N',
                            help='num decoder layers')
        parser.add_argument('--decoder-attention-heads', type=int, metavar='N',
                            help='num decoder attention heads')
        parser.add_argument('--decoder-normalize-before', default=False, action='store_true',
                            help='apply layernorm before each decoder block')
        parser.add_argument('--decoder-final-norm', default=False, action='store_true',
                            help='apply layernorm before each decoder block')
        parser.add_argument('--adaptive-softmax-cutoff', metavar='EXPR',
                            help='comma separated list of adaptive softmax cutoff points. '
                                 'Must be used with adaptive_loss criterion')
        parser.add_argument('--adaptive-softmax-dropout', type=float, metavar='D',
                            help='sets adaptive softmax dropout for the tail projections')
        parser.add_argument('--adaptive-softmax-factor', type=float, metavar='N',
                            help='adaptive input factor')
        parser.add_argument('--no-token-positional-embeddings', default=False, action='store_true',
                            help='if set, disables positional embeddings (outside self attention)')
        parser.add_argument('--share-decoder-input-output-embed', default=False, action='store_true',
                            help='share decoder input and output embeddings')
        parser.add_argument('--character-embeddings', default=False, action='store_true',
                            help='if set, uses character embedding convolutions to produce token embeddings')
        parser.add_argument('--character-filters', type=str, metavar='LIST',
                            default='[(1, 64), (2, 128), (3, 192), (4, 256), (5, 256), (6, 256), (7, 256)]',
                            help='size of character embeddings')
        parser.add_argument('--character-embedding-dim', type=int, metavar='N', default=4,
                            help='size of character embeddings')
        parser.add_argument('--char-embedder-highway-layers', type=int, metavar='N', default=2,
                            help='number of highway layers for character token embeddder')
        parser.add_argument('--adaptive-input', action='store_true',
                            help='if set, uses adaptive input')
        parser.add_argument('--adaptive-input-factor', type=float, metavar='N',
                            help='adaptive input factor')
        parser.add_argument('--adaptive-input-cutoff', metavar='EXPR',
                            help='comma separated list of adaptive input cutoff points.')
        parser.add_argument('--tie-adaptive-weights', action='store_true',
                            help='if set, ties the weights of adaptive softmax and adaptive input')
        parser.add_argument('--tie-adaptive-proj', action='store_true',
                            help='if set, ties the projection weights of adaptive softmax and adaptive input')
        parser.add_argument('--decoder-learned-pos', action='store_true',
                            help='use learned positional embeddings in the decoder')
        parser.add_argument('--decoder-embed-scale', default=True,
                            help='use learned positional embeddings in the decoder')
        parser.add_argument('--pretrained-checkpoint', metavar='DIR',
                            help='path to load checkpoint from pretrained model')
        parser.add_argument('--pretrained', default=False, action='store_true',
                            help='use pretrained model')
        # fmt: on

    @classmethod
    def build_model(cls, args, task):
        """Build a new model instance."""

        # make sure all arguments are present in older models
        base_lm_architecture_bpe(args)

        if not hasattr(args, 'max_source_positions'):
            args.max_source_positions = args.tokens_per_sample
        if not hasattr(args, 'max_target_positions'):
            args.max_target_positions = args.tokens_per_sample

        padding_idx = task.dictionary.pad() if task.dictionary.pad() != task.dictionary.eos() else None
        embed_tokens = Embedding(len(task.dictionary), args.decoder_input_dim, padding_idx)

        decoder = TransformerDecoder(
            args, task.output_dictionary, embed_tokens, no_encoder_attn=True,
            final_norm=args.decoder_final_norm,
        )
        return BPECDEFTransformerLanguageModel(decoder)


class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, dropout):
        super(MLP, self).__init__()

        layers = OrderedDict()
        layers['fc'] = nn.Linear(input_dim, hidden_dim)
        # layers['relu'] = nn.ReLU()
        layers['tanh'] = nn.Tanh()
        layers['drop'] = nn.Dropout(dropout)
        layers['out'] = nn.Linear(hidden_dim, 1)
        layers['sigmoid'] = nn.Sigmoid()

        self.model = nn.Sequential(layers)

    def forward(self, input):
        return self.model.forward(input)

class MLP_Discriminator(nn.Module):
    def __init__(self, embed_dim, output_dim):
        super(MLP_Discriminator, self).__init__()
        self.embed_dim = embed_dim
        self.hidden_dropout = 0.1
        self.input_dropout = 0.1

        self.mlp = MLP(embed_dim * 5, output_dim, self.hidden_dropout)
        self.dropout = nn.Dropout(self.input_dropout)

    def forward(self, s1, s2):
        inputs = torch.cat([s1, s2, s1 - s2, s1 * s2, torch.abs(s1 - s2)], -1)
        scores = self.mlp(self.dropout(inputs))
        return scores

class TransformerDecoder(FairseqIncrementalDecoder):
    """
    Transformer decoder consisting of *args.decoder_layers* layers. Each layer
    is a :class:`TransformerDecoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): decoding dictionary
        embed_tokens (torch.nn.Embedding): output embedding
        no_encoder_attn (bool, optional): whether to attend to encoder outputs
            (default: False).
        final_norm (bool, optional): apply layer norm to the output of the
            final decoder layer (default: True).
    """

    def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False, final_norm=True):
        super().__init__(dictionary)
        self.dropout = args.dropout
        self.share_input_output_embed = args.share_decoder_input_output_embed

        input_embed_dim = embed_tokens.embedding_dim
        embed_dim = args.decoder_embed_dim
        self.output_embed_dim = args.decoder_output_dim

        padding_idx = embed_tokens.padding_idx if embed_tokens.padding_idx is not None else None
        self.max_target_positions = args.max_target_positions

        self.sen_end_index = dictionary.enc.encode(' .')[0]
        self.pad_index = dictionary.pad()

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)  if args.decoder_embed_scale else None

        self.project_in_dim = Linear(input_embed_dim, embed_dim, bias=False) if embed_dim != input_embed_dim else None

        self.embed_positions = PositionalEmbedding(
            args.max_target_positions, embed_dim, padding_idx,
            learned=args.decoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerDecoderLayer(args, no_encoder_attn)
            for _ in range(args.decoder_layers)
        ])

        self.register_buffer('version', torch.Tensor([2]))
        self.normalize = args.decoder_normalize_before and final_norm
        if self.normalize:
            self.layer_norm = LayerNorm(embed_dim)

        # only train discriminator layer
        for p in self.parameters():
            p.requires_grad = False
        self.discriminator = MLP_Discriminator(input_embed_dim, self.output_embed_dim)

    def forward(self, prev_output_tokens, encoder_out=None, incremental_state=None, **unused):
        """
        Args:
            prev_output_tokens (LongTensor): previous decoder outputs of shape
                `(batch, tgt_len)`, for input feeding/teacher forcing
            encoder_out (Tensor, optional): output from the encoder, used for
                encoder-side attention
            incremental_state (dict): dictionary used for storing state during
                :ref:`Incremental decoding`

        Returns:
            tuple:
                - the decoder's output of shape `(batch, tgt_len, vocab)`
                - a dictionary with any model-specific outputs
        """
        pos_result = self.extract_features(prev_output_tokens, encoder_out, incremental_state)
        return pos_result

    def encoder_text(self, prev_output_tokens, encoder_out=None, incremental_state=None):
        # embed positions
        positions = self.embed_positions(
            prev_output_tokens,
            incremental_state=incremental_state,
        ) if self.embed_positions is not None else None

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:]
            if positions is not None:
                positions = positions[:, -1:]

        # embed tokens and positions
        x = self.embed_tokens(prev_output_tokens)
        if self.embed_scale is not None:
            x *= self.embed_scale

        if self.project_in_dim is not None:
            x = self.project_in_dim(x)

        if positions is not None:
            x += positions
        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        inner_states = [x]

        # decoder layers
        for layer in self.layers:
            x, attn = layer(
                x,
                encoder_out['encoder_out'] if encoder_out is not None else None,
                encoder_out['encoder_padding_mask'] if encoder_out is not None else None,
                incremental_state,
                self_attn_mask=self.buffered_future_mask(x) if incremental_state is None else None,
            )
            inner_states.append(x)

        if self.normalize:
            x = self.layer_norm(x)

        # T x B x C -> B x T x C
        x = x.transpose(0, 1)

        return x

    def get_sentence_state(self, hidden_state, word_sen_label, max_sen_len):
        x = hidden_state
        bsz = hidden_state.size(0)
        # # extract sentence state
        sentence_state = x.new_full((bsz, max_sen_len, self.output_embed_dim), 0, dtype=torch.float)
        for i in range(max_sen_len):
            cur_sen_mask = (word_sen_label == i)
            cur_sen_len = cur_sen_mask.sum(dim=1)
            cur_sen_len[cur_sen_len==0] = 1
            cur_sen_state = (cur_sen_mask.float().unsqueeze(dim=-1) * x).sum(dim=1) / cur_sen_len.float().unsqueeze(dim=-1)

            sentence_state[:, i] = cur_sen_state
        return sentence_state

    def get_sentence_label(self, prev_output_tokens):
        bsz, time = prev_output_tokens.size()
        sen_sympol_mask = (prev_output_tokens == self.sen_end_index)
        prev_mask = (prev_output_tokens == self.pad_index)
        all_inds = torch.arange(time).cuda()

        word_sen_label = prev_output_tokens.new_full((bsz, time), -1, dtype=torch.long)
        max_sen_len = 0
        # for each story, for each sentence, get sentence state and give sentence label
        for i in range(bsz):
            # story
            if (sen_sympol_mask[i]).any() and (sen_sympol_mask[i]).nonzero().size(0) > 3:
                sens = (sen_sympol_mask[i]).nonzero() + 1
                sen_len = min(100, sens.size(0))
                max_sen_len = sen_len if sen_len > max_sen_len else max_sen_len
                for j in range(sen_len):
                    # exclude <s>
                    last = 1 if j == 0 else sens[j - 1]
                    word_sen_label[i, last: sens[j]] = j
                # word_sen_label[i, sens[-1]:] = sen_len
            else:
                segments_length = 15
                segments_x = torch.split(prev_output_tokens[i], segments_length, dim=0)
                sen_len = min(100, len(segments_x))
                max_sen_len = sen_len if sen_len > max_sen_len else max_sen_len
                for ii, seg in enumerate(segments_x):
                    cur_word_mask = prev_mask[i, ii * segments_length: (ii + 1) * segments_length]
                    if cur_word_mask.any():
                        if cur_word_mask.all():
                            word_sen_label[i, ii * segments_length: (ii + 1) * segments_length] =  -1
                        else:
                            cur_ind_mask = ( all_inds >= (ii * segments_length)) & (all_inds < ((ii + 1) * segments_length))
                            part_word_ind = (~prev_mask[i]) & cur_ind_mask
                            word_sen_label[i, part_word_ind] = ii
                    else:
                        word_sen_label[i, ii * segments_length: (ii + 1) * segments_length] = ii

        sentence_length = word_sen_label.max(dim=1)[0] + 1
        return word_sen_label, max_sen_len, sentence_length

    def get_pair_score(self, tokens, encoder_out, incremental_state, word_sen_label, max_sen_len):
        # gen sentence state
        hidden_state = self.encoder_text(tokens, encoder_out, incremental_state)
        sentence_state = self.get_sentence_state(hidden_state, word_sen_label, max_sen_len)

        first = sentence_state[:, :-1, :]
        second = sentence_state[:, 1:, :]
        score = self.discriminator(first, second)
        return score, first, second

    def extract_features(self, prev_output_tokens, encoder_out=None, incremental_state=None, **unused):
        """
        Similar to *forward* but only return features.

        Returns:
            tuple:
                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
                - a dictionary with any model-specific outputs
        """

        # get sen split
        bsz, time = prev_output_tokens.size()
        word_sen_label, max_sen_len, sentence_length = self.get_sentence_label(prev_output_tokens)

        # positive sentence pair, sentence state: b x s x d, score: b x s x 1
        pos_scores, sent1, pos_sent2  = self.get_pair_score(prev_output_tokens, encoder_out, incremental_state, word_sen_label, max_sen_len)
        pos_scores = pos_scores.squeeze(dim=2)

        # get mask (b x (s-1), b)
        score_pad_mask = (pos_sent2.sum(dim=-1) != 0)

        # article score
        temp_pos_scores = pos_scores.clone()
        temp_pos_scores[~score_pad_mask] = 0
        pos_article_scores = temp_pos_scores.sum(dim=1)
        pos_article_lens = score_pad_mask.sum(dim=-1)
        pos_article_lens[pos_article_lens == 0] = 1
        pos_article_scores = pos_article_scores / pos_article_lens.float()

        # 1. negative sentence pair, easy negative examples, shuffle
        neg_shuffle_tokens = torch.ones_like(prev_output_tokens) * self.pad_index
        shuffle_word_sen_label = torch.ones_like(prev_output_tokens) * -1
        for i in range(bsz):
            shuffle_order = torch.randperm(sentence_length[i]).tolist()
            start_ind = 1
            new_label = 0
            for sen_label in shuffle_order:
                cur_sen = word_sen_label[i] == sen_label
                end_ind = start_ind + cur_sen.nonzero().size(0)
                neg_shuffle_tokens[i, start_ind: end_ind] = prev_output_tokens[i, cur_sen]
                shuffle_word_sen_label[i, start_ind: end_ind] = new_label

                new_label += 1
                start_ind = end_ind

        neg_scores_shuffle, _, _ = self.get_pair_score(neg_shuffle_tokens, encoder_out, incremental_state, shuffle_word_sen_label, max_sen_len)
        # neg_scores_shuffle = neg_scores_shuffle.squeeze(dim=2)

        # 3. harder examples, left shift
        neg_swap_tokens = torch.zeros_like(prev_output_tokens) * self.pad_index
        swap_word_sen_label = torch.ones_like(prev_output_tokens) * -1
        for i in range(bsz):
            shuffle_order = range(sentence_length[i].item())[::-1]
            start_ind = 1
            new_label = 0
            for sen_label in shuffle_order:
                cur_sen = word_sen_label[i] == sen_label
                end_ind = start_ind + cur_sen.nonzero().size(0)
                neg_swap_tokens[i, start_ind: end_ind] = prev_output_tokens[i, cur_sen]
                swap_word_sen_label[i, start_ind: end_ind] = new_label

                new_label += 1
                start_ind = end_ind

        neg_scores_swap, _, _ = self.get_pair_score(neg_swap_tokens, encoder_out, incremental_state, swap_word_sen_label, max_sen_len)
        # neg_scores_shift_left = neg_scores_shift_left.squeeze(dim=2)

        # 4. repeat pair
        # neg_scores_repeat = self.discriminator(sent1, sent1)
        # neg_scores_repeat = neg_scores_repeat.squeeze(dim=2)

        # weight sum
        neg_scores = torch.cat([neg_scores_shuffle, neg_scores_swap], dim=-1)
        # neg_weight = F.softmax(neg_scores * 2, dim=-1)
        # neg_scores = (neg_scores * neg_weight).sum(dim=-1)
        neg_scores = neg_scores.mean(dim=-1)

        temp_pos_scores = neg_scores.clone()
        temp_pos_scores[~score_pad_mask] = 0
        neg_article_scores = temp_pos_scores.sum(dim=1)
        neg_article_scores = neg_article_scores / pos_article_lens.float()

        article_pair_acc = (pos_article_scores > neg_article_scores).float().mean()

        # sentence level word reward
        # ignore pad word
        word_sen_label[word_sen_label==-1] = 0
        return_word_scores = pos_scores
        first_sen_score = return_word_scores.new_ones((bsz, 1))
        temp_pos_scores = torch.cat([first_sen_score, return_word_scores], dim=-1)
        temp_pos_scores = temp_pos_scores.unsqueeze(dim=1).expand(-1, time, -1)
        word_scores = temp_pos_scores.contiguous().view(-1, temp_pos_scores.size(-1)).gather(dim=-1, index=word_sen_label.view(-1, 1))

        # for debug
        sentence_pair_acc = (pos_scores[score_pad_mask] > neg_scores[score_pad_mask]).float().mean()
        neg_1 = neg_scores_shuffle[score_pad_mask].mean()
        neg_2 = neg_scores_swap[score_pad_mask].mean()


        return_scores = {'pos_scores': pos_scores, 'neg_scores': neg_scores, 'score_pad_mask': score_pad_mask,
                         'pos_article_scores': pos_article_scores, 'word_scores': word_scores,
                         'neg_article_scores': neg_article_scores, 'article_pair_acc': article_pair_acc,
                         'sentence_pair_acc': sentence_pair_acc,
                         'neg_1': neg_1, 'neg_2': neg_2,
                         }
        return return_scores


    def max_positions(self):
        """Maximum output length supported by the decoder."""
        if self.embed_positions is None:
            return self.max_target_positions
        return min(self.max_target_positions, self.embed_positions.max_positions())

    def buffered_future_mask(self, tensor):
        dim = tensor.size(0)
        if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device:
            self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
        if self._future_mask.size(0) < dim:
            self._future_mask = torch.triu(utils.fill_with_neg_inf(self._future_mask.resize_(dim, dim)), 1)
        return self._future_mask[:dim, :dim]

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)

        for i in range(len(self.layers)):
            # update layer norms
            layer_norm_map = {
                '0': 'self_attn_layer_norm',
                '1': 'encoder_attn_layer_norm',
                '2': 'final_layer_norm'
            }
            for old, new in layer_norm_map.items():
                for m in ('weight', 'bias'):
                    k = '{}.layers.{}.layer_norms.{}.{}'.format(name, i, old, m)
                    if k in state_dict:
                        state_dict['{}.layers.{}.{}.{}'.format(name, i, new, m)] = state_dict[k]
                        del state_dict[k]
        if utils.item(state_dict.get('{}.version'.format(name), torch.Tensor([1]))[0]) < 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict['{}.version'.format(name)] = torch.Tensor([1])

        return state_dict

class TransformerDecoderLayer(nn.Module):
    """Decoder layer block.

    In the original paper each operation (multi-head attention, encoder
    attention or FFN) is postprocessed with: `dropout -> add residual ->
    layernorm`. In the tensor2tensor code they suggest that learning is more
    robust when preprocessing each layer with layernorm and postprocessing with:
    `dropout -> add residual`. We default to the approach in the paper, but the
    tensor2tensor approach can be enabled by setting
    *args.decoder_normalize_before* to ``True``.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        no_encoder_attn (bool, optional): whether to attend to encoder outputs
            (default: False).
    """

    def __init__(self, args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False):
        super().__init__()
        self.embed_dim = args.decoder_embed_dim
        self.self_attn = MultiheadAttention(
            embed_dim=self.embed_dim,
            num_heads=args.decoder_attention_heads,
            dropout=args.attention_dropout,
            add_bias_kv=add_bias_kv,
            add_zero_attn=add_zero_attn,
        )
        self.dropout = args.dropout
        self.activation_fn = utils.get_activation_fn(
            activation=getattr(args, 'activation_fn', 'relu')
        )
        self.activation_dropout = getattr(args, 'activation_dropout', 0)
        if self.activation_dropout == 0:
            # for backwards compatibility with models that use args.relu_dropout
            self.activation_dropout = getattr(args, 'relu_dropout', 0)
        self.normalize_before = args.decoder_normalize_before

        # use layerNorm rather than FusedLayerNorm for exporting.
        # char_inputs can be used to determint this.
        # TODO  remove this once we update apex with the fix
        export = getattr(args, 'char_inputs', False)
        self.self_attn_layer_norm = LayerNorm(self.embed_dim, export=export)

        if no_encoder_attn:
            self.encoder_attn = None
            self.encoder_attn_layer_norm = None
        else:
            self.encoder_attn = MultiheadAttention(
                self.embed_dim, args.decoder_attention_heads,
                dropout=args.attention_dropout,
            )
            self.encoder_attn_layer_norm = LayerNorm(self.embed_dim, export=export)

        self.fc1 = Linear(self.embed_dim, args.decoder_ffn_embed_dim)
        self.fc2 = Linear(args.decoder_ffn_embed_dim, self.embed_dim)

        self.final_layer_norm = LayerNorm(self.embed_dim, export=export)
        self.need_attn = True

        self.onnx_trace = False

    def prepare_for_onnx_export_(self):
        self.onnx_trace = True

    def forward(
        self,
        x,
        encoder_out=None,
        encoder_padding_mask=None,
        incremental_state=None,
        prev_self_attn_state=None,
        prev_attn_state=None,
        self_attn_mask=None,
        self_attn_padding_mask=None,
    ):
        """
        Args:
            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
            encoder_padding_mask (ByteTensor): binary ByteTensor of shape
                `(batch, src_len)` where padding elements are indicated by ``1``.

        Returns:
            encoded output of shape `(batch, src_len, embed_dim)`
        """
        residual = x
        x = self.maybe_layer_norm(self.self_attn_layer_norm, x, before=True)
        if prev_self_attn_state is not None:
            if incremental_state is None:
                incremental_state = {}
            prev_key, prev_value = prev_self_attn_state
            saved_state = {"prev_key": prev_key, "prev_value": prev_value}
            self.self_attn._set_input_buffer(incremental_state, saved_state)
        x, attn = self.self_attn(
            query=x,
            key=x,
            value=x,
            key_padding_mask=self_attn_padding_mask,
            incremental_state=incremental_state,
            need_weights=False,
            attn_mask=self_attn_mask,
        )
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.maybe_layer_norm(self.self_attn_layer_norm, x, after=True)

        if self.encoder_attn is not None:
            residual = x
            x = self.maybe_layer_norm(self.encoder_attn_layer_norm, x, before=True)
            if prev_attn_state is not None:
                if incremental_state is None:
                    incremental_state = {}
                prev_key, prev_value = prev_attn_state
                saved_state = {"prev_key": prev_key, "prev_value": prev_value}
                self.encoder_attn._set_input_buffer(incremental_state, saved_state)
            x, attn = self.encoder_attn(
                query=x,
                key=encoder_out,
                value=encoder_out,
                key_padding_mask=encoder_padding_mask,
                incremental_state=incremental_state,
                static_kv=True,
                need_weights=(not self.training and self.need_attn),
            )
            x = F.dropout(x, p=self.dropout, training=self.training)
            x = residual + x
            x = self.maybe_layer_norm(self.encoder_attn_layer_norm, x, after=True)

        residual = x
        x = self.maybe_layer_norm(self.final_layer_norm, x, before=True)
        x = self.activation_fn(self.fc1(x))
        x = F.dropout(x, p=self.activation_dropout, training=self.training)
        x = self.fc2(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
        x = self.maybe_layer_norm(self.final_layer_norm, x, after=True)
        if self.onnx_trace and incremental_state is not None:
            saved_state = self.self_attn._get_input_buffer(incremental_state)
            self_attn_state = saved_state["prev_key"], saved_state["prev_value"]
            return x, attn, self_attn_state
        return x, attn

    def maybe_layer_norm(self, layer_norm, x, before=False, after=False):
        assert before ^ after
        if after ^ self.normalize_before:
            return layer_norm(x)
        else:
            return x

    def make_generation_fast_(self, need_attn=False, **kwargs):
        self.need_attn = need_attn


def Embedding(num_embeddings, embedding_dim, padding_idx):
    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
    nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
    if padding_idx is not None:
        nn.init.constant_(m.weight[padding_idx], 0)
    return m


def Linear(in_features, out_features, bias=True):
    m = nn.Linear(in_features, out_features, bias)
    nn.init.xavier_uniform_(m.weight)
    if bias:
        nn.init.constant_(m.bias, 0.)
    return m

@register_model_architecture('transformer_lm_bpe_coherence_disef', 'transformer_lm_bpe_coherence_disef')
def base_lm_architecture_bpe(args):
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 512)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 2048)
    args.decoder_layers = getattr(args, 'decoder_layers', 6)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 8)
    args.adaptive_softmax_cutoff = getattr(args, 'adaptive_softmax_cutoff', None)
    args.adaptive_softmax_dropout = getattr(args, 'adaptive_softmax_dropout', 0)
    args.adaptive_softmax_factor = getattr(args, 'adaptive_softmax_factor', 4)
    args.decoder_learned_pos = getattr(args, 'decoder_learned_pos', False)
    args.activation_fn = getattr(args, 'activation_fn', 'relu')

    args.add_bos_token = getattr(args, 'add_bos_token', False)
    args.character_embeddings = getattr(args, 'character_embeddings', False)

    args.decoder_output_dim = getattr(args, 'decoder_output_dim', args.decoder_embed_dim)
    args.decoder_input_dim = getattr(args, 'decoder_input_dim', args.decoder_embed_dim)

    # The model training is not stable without this
    args.decoder_normalize_before = True

    args.adaptive_input = getattr(args, 'adaptive_input', False)
    args.adaptive_input_factor = getattr(args, 'adaptive_input_factor', 4)
    args.adaptive_input_cutoff = getattr(args, 'adaptive_input_cutoff', None)

    args.tie_adaptive_weights = getattr(args, 'tie_adaptive_weights', False)
    args.tie_adaptive_proj = getattr(args, 'tie_adaptive_proj', False)

@register_model_architecture('transformer_lm_bpe_coherence_disef', 'transformer_lm_bpe_coherence_disef_gpt_format')
def transformer_lm_gpt_bpe_format(args):
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
    args.decoder_layers = getattr(args, 'decoder_layers', 12)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 12)
    args.dropout = getattr(args, 'dropout', 0.1)
    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
    args.activation_fn = getattr(args, 'activation_fn', 'gelu_accurate')
    args.share_decoder_input_output_embed = True
    args.decoder_normalize_before = True
    args.decoder_final_norm = True
    args.decoder_learned_pos = True
    args.decoder_embed_scale = False
    base_lm_architecture_bpe(args)
